# -*- coding: utf-8 -*-
"""
@ Created on 2024-06-14 12:23
---------
@summary: 数据库模块
---------
@author: Lsaiah
"""

from tortoise import Tortoise
from settings import settings
from log import logger


# TortoiseORM 配置
TORTOISE_ORM = {
    "connections": {
        "default": {
            "engine": "tortoise.backends.mysql",
            "credentials": {
                "host": settings.MYSQL_IP,
                "port": settings.MYSQL_PORT,
                "user": settings.MYSQL_USER_NAME,
                "password": settings.MYSQL_USER_PASS,
                "database": settings.MYSQL_DB,
                "charset": "utf8mb4",
            }
        }
    },
    "apps": {
        "models": {
            "models": [
                "app.system.user.models.curd.user_crud",
                "app.system.dict.models.curd.dict_type_crud",
                "app.system.dict.models.curd.dict_item_crud",
                "aerich.models"
            ],
            "default_connection": "default",
        }
    },
    "use_tz": False,
    "timezone": "Asia/Shanghai"
}


async def init_db():
    """初始化数据库连接"""
    try:
        await Tortoise.init(config=TORTOISE_ORM)
        logger.info("数据库连接初始化成功")
        
        # 自动生成表结构（如果不存在）
        try:
            await Tortoise.generate_schemas(safe=True)
            logger.info("数据库表结构检查/生成完成")
        except Exception as e:
            logger.warning(f"数据库表结构生成失败: {str(e)}")
        
    except Exception as e:
        logger.error(f"数据库初始化失败: {str(e)}")
        raise


async def close_db():
    """关闭数据库连接"""
    try:
        await Tortoise.close_connections()
        logger.info("数据库连接已关闭")
    except Exception as e:
        logger.error(f"关闭数据库连接失败: {str(e)}")


async def create_tables():
    """手动创建数据库表（用于初始化）"""
    try:
        await Tortoise.init(config=TORTOISE_ORM)
        await Tortoise.generate_schemas(safe=True)
        logger.info("数据库表创建/检查完成")
        await Tortoise.close_connections()
    except Exception as e:
        logger.error(f"数据库表创建失败: {str(e)}")
        raise


async def get_db():
    """获取数据库连接"""
    return Tortoise.get_connection("default")


# 数据库事务装饰器
async def with_transaction(func):
    """数据库事务装饰器"""
    async def wrapper(*args, **kwargs):
        conn = Tortoise.get_connection("default")
        async with conn.in_transaction():
            return await func(*args, **kwargs)
    return wrapper


# 数据库连接健康检查
async def check_db_health():
    """检查数据库连接健康状态"""
    try:
        conn = Tortoise.get_connection("default")
        await conn.execute_query("SELECT 1")
        return True
    except Exception as e:
        logger.error(f"数据库健康检查失败: {str(e)}")
        return False


__all__ = [
    "TORTOISE_ORM", "init_db", "close_db", "get_db",
    "with_transaction", "check_db_health", "create_tables"
]
